#!/bin/bash

function showhelp {
    echo 'usage : cp2ks option dir_contain_cp2kinp or cp2k.inp'
    echo 'e.g.  : cp2ks -n 128 -N 2  -e node05 ABC_OPT'
    echo 'option: -[n total number of cores, a for all core with optimal settings]'
    echo '        -[N num of openmp thread, number of rank is n/N]'
    echo '        -[l node name. e.g. -l node01,node02] '
    echo '        -[e nodename to exclude. e.g. -e node01,node02]'
    echo '        -[q queue name to submit]'
    echo '        -[E exclusive mode]'
    echo '        -[p pbs or slurm]'
    echo '        -[m map setting for mpirun, n for not use mapping, d(default) for default setting]'
    echo '        -[R using the restart file]'
}
if [ $# -lt 1 ]; then
    showhelp
    exit 1
fi
while getopts ":n:N:t:l:e:q:ERm:" opt
do
    case $opt in
        n) NPROC=$OPTARG;;
        N) THREAD=$OPTARG;;
        l) NODE=$OPTARG;;
        e) EXCLUDE=$OPTARG;;
        q) QUEUE=$OPTARG;;
        m) MAPPING=$OPTARG;;
        E) EXCLUSIVE=1;;
        R) RESTART=1;;
        ?) showhelp
        exit 1;;
    esac
done
shift $(($OPTIND - 1))

if [ -z "$NPROC" ];then
    NPROC=64
fi


if [[ "$NPROC" != "a" ]];then
    if [ -z "$THREAD" ];then
            THREAD=1
    fi
    RANK=`echo "$NPROC/$THREAD" | bc`
else
    THREAD=""
    RANK=""
    EXCLUSIVE=1
fi

if [ -z "$MAPPING" ];then
    MAPPING=d
fi

function check_cpu() {
   cpucore=`lscpu | awk '{if($1 ~/Core/)print $4}'`
   socket=`lscpu | awk '{if($1 ~/Socket/)print $2}'`
   L3=`lscpu | awk '{if($1 ~/L3/){gsub(/\(/,"",$5);print $5}}'`
   L3core=`awk -v c=$cpucore -v s=$socket -v l=$L3 'BEGINN{print c*s/l}'`
   allcore=`awk -v c=$cpucore -v s=$socket 'BEGINN{print c*s}'`
}


while [[ $# -ge 1 ]]
do

# check job manage system
PBS_FLAG=`which qsub 2> null`
SLURM_FLAG=`which sbatch 2> null`
if [ -n "$SLURM_FLAG" ];then
    JOB_NAME=${1%.*}
    NAME=sub_${JOB_NAME}.sbatch
    INSUBDIR=''
    if [ -d "$1" ];then
        cd $1
        pwd
        INSUBDIR=1
        if [[ $RESTART == 1 ]];then
           inpname=`ls *.restart`
        else
            inpname=`ls *.inp`
        fi
    else
        inpname=$1
    fi
    echo '#!/bin/bash' > $NAME
    echo "#SBATCH -J $JOB_NAME" >> $NAME
    echo "#SBATCH -N 1" >> $NAME
    if [ -n "$NODE" ];then
        echo "#SBATCH -w $NODE" >> $NAME
    fi
    if [ -n "$EXCLUDE" ];then
        echo "#SBATCH --exclude $EXCLUDE" >> $NAME
    fi
    if [ -n "$EXCLUSIVE" ];then
        echo "#SBATCH --exclusive" >> $NAME
    fi
    if [ -n "$QUEUE" ];then
        echo "#SBATCH --partition=$QUEUE" >> $NAME
    fi
    if [[ -n "$RANK" ]];then
        echo "#SBATCH --ntasks-per-node=$RANK" >> $NAME
    fi
    if [[ -n "$THREAD" ]];then
        echo "#SBATCH --cpus-per-task=$THREAD" >> $NAME
    fi
    if [ -n "$NGPUS" ];then
        echo "#SBATCH -G $NGPUS" >> $NAME
    fi
    echo "source ~/.bashrc" >> $NAME

    if [[ $RANK == 144 ]] && [[ $MAPPING == "d" ]];then
        echo 'map="-map-by ppr:6:l3cache:pe=1"' >> $NAME
        echo 'RANK=144' >> $NAME
        echo 'THREAD=1' >> $NAME
    elif [[ $RANK == 64 ]] && [[ $THREAD == 2 ]] && [[ $MAPPING == "d" ]];then
        echo 'map="-map-by ppr:4:l3cache:pe=2"' >> $NAME
        echo 'RANK=64' >> $NAME
        echo 'THREAD=2' >> $NAME
    fi
    if [[ $NPROC == "a" ]];then
        echo "cpucore=\`lscpu | awk '{if(\$1 ~/Core/)print \$4}'\`" >> $NAME
        echo "socket=\`lscpu | awk '{if(\$1 ~/Socket/)print \$2}'\`" >> $NAME
        echo "L3=\`lscpu | awk '{if(\$1 ~/L3/){gsub(/\(/,\"\",\$5);print \$5}}'\`" >> $NAME
        echo "L3core=\`awk -v c=\$cpucore -v s=\$socket -v l=\$L3 'BEGIN{print c*s/l}'\`" >> $NAME
        echo "allcore=\`awk -v c=\$cpucore -v s=\$socket 'BEGIN{print c*s}'\`" >> $NAME
        echo 'if [[ $allcore == 192 ]];then' >> $NAME
        echo '    RANK=96' >> $NAME
        echo '    THREAD=2' >> $NAME
        echo '    map="-map-by ppr:4:l3cache:pe=2"' >> $NAME
        echo 'elif [[ $allcore == 128 ]];then' >> $NAME
        echo '    RANK=64' >> $NAME
        echo '    THREAD=2' >> $NAME
        echo '    map="-map-by ppr:4:l3cache:pe=2"' >> $NAME
        echo 'fi' >> $NAME
    else
        echo "    RANK=$RANK" >> $NAME
        echo "    THREAD=$THREAD" >> $NAME
    fi
    echo "export OMP_NUM_THREADS=\$THREAD" >> $NAME
    echo "export OMP_STACKSIZE=512m"  >> $NAME
    echo 'start=`date +%s`' >> $NAME
    echo 'module load cp2k/2023.2-openmpi416-gcc113' >> $NAME
    echo "mpirun -np \$RANK  \$map  cp2k.psmp -o ${inpname%.*}.out $inpname " >> $NAME
    echo 'end=`date +%s`' >> $NAME
    echo 'runtime=$(echo "$end - $start" | bc -l )' >> $NAME
    echo "echo \$runtime > runtime_r${RANK}_n${THREAD}" >> $NAME
    sbatch $NAME
    if [ -n "$INSUBDIR" ];then
        cd ..
    fi
elif [ -n "$PBS_FLAG" ];then
    JOB_NAME=${1/.sh}
    NAME=sub_${1/.sh/}.pbs
    echo "#PBS -S /bin/bash" > $NAME
    if [ -n "$NODE" ];then
        echo "#PBS -l nodes=$NODE:ppn=$NPROC" >> $NAME
    else
        echo "#PBS -l nodes=1:ppn=$NPROC" >> $NAME
    fi
    echo "#PBS -N $JOB_NAME" >> $NAME
    echo "#PBS -j oe" >> $NAME
    echo 'cd $PBS_O_WORKDIR' >> $NAME
    echo "source ~/.bashrc" >> $NAME
    if [ -d "${1}" ]; then
        echo "cd $1" >> $NAME
        echo 'sh *.sh' >> $NAME
    else
        echo "sh $1 &> ${1%.*}.out" >> $NAME
    fi
    qsub $NAME
fi
shift 1
done 
